package org.apache.flink.runtime.minicluster;

import java.util.concurrent.CompletableFuture;

import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.api.common.ConsumingStatus;
import org.apache.flink.api.common.ResourceStatus;
import org.apache.flink.core.execution.JobClient;
import org.apache.flink.runtime.allocation.ResourceAllocationStatistics;
import org.apache.flink.runtime.client.ConsumeClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** A {@link JobClient} for a {@link MiniCluster}. */
public final class MiniClusterConsumerClient implements ConsumeClient {

    private static final Logger LOG = LoggerFactory.getLogger(MiniClusterConsumerClient.class);

    private final ConsumerID consumerId;
    private final MiniCluster miniCluster;
    private final ClassLoader classLoader;
    private final CompletableFuture<ResourceAllocationStatistics> consumingFuture;

    /**
     * Creates a {@link MiniClusterConsumerClient} for the given {@link ConsumerID} and {@link MiniCluster}.
     * This will shut down the {@code MiniCluster} after job result retrieval if {@code
     * shutdownCluster} is {@code true}.
     */
    public MiniClusterConsumerClient(
            ConsumerID consumerId,
            MiniCluster miniCluster,
            ClassLoader classLoader,
            ConsumeFinalizationBehavior finalizationBehaviour) {
        this.consumerId = consumerId;
        this.miniCluster = miniCluster;
        this.classLoader = classLoader;

        if (finalizationBehaviour == ConsumeFinalizationBehavior.SHUTDOWN_CLUSTER) {
            // Make sure to shutdown the cluster when the job completes.
            consumingFuture = miniCluster
                    		.requestResourceAllocationStatistics()
                            .whenComplete((result, throwable) -> shutDownCluster(miniCluster));
        } else if (finalizationBehaviour == ConsumeFinalizationBehavior.NOTHING) {
            consumingFuture = miniCluster.requestResourceAllocationStatistics();
        } else {
            throw new IllegalArgumentException(
                    "Unexpected shutdown behavior: " + finalizationBehaviour);
        }
    }

    @Override
    public ConsumerID getConsumerID() {
        return consumerId;
    }

    @Override
    public CompletableFuture<ResourceAllocationStatistics> getResourceAllocationStatistics() {
        return miniCluster.requestResourceAllocationStatistics();
    }

    @Override
    public CompletableFuture<Void> cancel() {
        return miniCluster.cancel().thenAccept(result -> {});
    }
   
    @Override
	public CompletableFuture<ConsumingStatus> getConsumingStatus() {
    	return miniCluster.getConsumingStatus();
	}

	@Override
	public CompletableFuture<ResourceStatus> getResourceStatus() {
		return miniCluster.getResourceStatus();
	}

	private static void shutDownCluster(MiniCluster miniCluster) {
        miniCluster
                .closeAsync()
                .whenComplete(
                        (ignored, throwable) -> {
                            if (throwable != null) {
                                LOG.warn("Shutdown of MiniCluster failed.", throwable);
                            }
                        });
    }

    /** Determines the behavior of the {@link MiniClusterConsumerClient} when the consuming finishes. */
    public enum ConsumeFinalizationBehavior {
        /** Shut down the {@link MiniCluster} when the job finishes. */
        SHUTDOWN_CLUSTER,

        /** Don't do anything when the job finishes. */
        NOTHING
    }
}
